#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
CodeSearchNet-Python 评测脚本
在处理后的数据集上测试代码生成能力，使用 AST 结构和关键词匹配进行初步评估
输出：通过率和准确率统计
"""
import os
import json
import ast
import sys
from typing import List, Dict, Optional

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.inference import generate_code


def _first_existing(paths: List[str]) -> Optional[str]:
    """返回路径列表中第一个存在的路径"""
    for p in paths:
        if os.path.exists(p):
            return p
    return None


def _detect_codesearchnet_paths():
    """检测处理后的 CodeSearchNet-Python 测试集路径
    
    Returns:
        (test_path, test_candidates): 测试集路径和候选路径列表
    """
    test_candidates = [
        os.path.join('data', 'processed', 'codesearchnet_test_text2code.jsonl'),  # 新命名
        os.path.join('data', 'processed', 'complex_test_text2code.jsonl'),        # 当前命名
    ]
    return _first_existing(test_candidates), test_candidates


def _read_jsonl(path: str) -> List[Dict]:
    """读取 JSONL 格式文件"""
    data = []
    if not os.path.exists(path):
        return data
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                data.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return data


def _is_valid_python_function(code: str) -> bool:
    """检查代码是否为语法正确的 Python 函数"""
    s = (code or '').strip()
    if not s.startswith('def '):
        return False
    try:
        tree = ast.parse(s)
    except SyntaxError:
        return False
    fdefs = [n for n in tree.body if isinstance(n, ast.FunctionDef)]
    return len(fdefs) >= 1


def _rough_quality_check(prompt: str, code: str) -> bool:
    """基于关键词和 AST 结构进行粗略质量检查"""
    s = (code or '').strip().lower()
    p = (prompt or '').strip().lower()
    
    if not _is_valid_python_function(code):
        return False
    
    # 排除明显错误的模式
    if 'class ' in s:
        # 要求函数但生成了类，视为失败
        if 'function' in p and 'class' not in p:
            return False
    
    # 关键词匹配检查
    if 'return' not in s:
        # 大多数函数应有返回值
        if 'print' not in p:  # 除非是打印函数
            return False
    
    return True


def main(n: int | None = None, model: str = "auto"):
    """主评测函数
    
    Args:
        n: 评测样本数量，None 表示全部
        model: 模型类型 (auto/complex/simple/multitask)
    """
    # 根据选择设置模型环境变量
    if model == "multitask":
        os.environ["T2C_INFER_MODE"] = "multitask"
        print("[INFO] 评测多任务 LoRA 模型")
    elif model == "complex":
        os.environ["T2C_INFER_MODE"] = "complex"
        print("[INFO] 评测复杂任务 LoRA 模型")
    elif model == "simple":
        os.environ["T2C_INFER_MODE"] = "simple"
        print("[INFO] 评测简单任务 LoRA 模型")
    else:  # auto
        os.environ["T2C_INFER_MODE"] = "auto"
        print("[INFO] 自动检测最佳可用模型")
    
    test_path, test_candidates = _detect_codesearchnet_paths()
    test_records = _read_jsonl(test_path) if test_path else []
    
    if not test_records:
        looked = ", ".join(test_candidates)
        print(f"[ERROR] 未找到处理后的 CodeSearchNet-Python 测试集。查找路径: {looked}\n"
              f"请先运行预处理脚本 (如 preprocess_data.py) 生成数据集\n"
              f"提示: 当前复杂任务数据集保存为 'data/processed/complex_*_text2code.jsonl'")
        return 1

    total = 0
    pass_cnt = 0
    
    for item in (test_records[:n] if n else test_records):
        instr = item.get('instruction', '')
        total += 1
        
        # 使用 complex 模式生成代码
        code = generate_code(instr, force_few_shots=True, mode="complex")
        
        # 使用粗略规则检查
        ok = _rough_quality_check(instr, code)
        
        if ok:
            pass_cnt += 1
        
        print(f"[CASE {total}] {'PASS' if ok else 'FAIL'}")

    acc = pass_cnt / total if total else 0.0
    print(f"[RESULT] 通过 {pass_cnt}/{total}  (准确率={acc:.3f})")
    return 0


if __name__ == '__main__':
    import sys
    import argparse
    parser = argparse.ArgumentParser(description="CodeSearchNet-Python 评测脚本")
    parser.add_argument('--n', type=int, default=None, help="评测样本数量 (默认: 全部)")
    parser.add_argument('--model', type=str, default="auto", 
                        choices=["auto", "complex", "simple", "multitask"],
                        help="评测模型类型 (默认: auto)")
    args = parser.parse_args()
    raise SystemExit(main(args.n, args.model))
